﻿import cv2
import numpy as np
import matplotlib.pyplot as plt

def embed_watermark(img_path, watermark_text="SDUT 计算机"):
    # 读取图像并转为YUV格式（保留Y通道处理）
    img = cv2.imread(img_path)
    yuv = cv2.cvtColor(img, cv2.COLOR_BGR2YUV)
    Y = yuv[:,:,0].astype(np.float32)
    
    # 1. DFT变换
    dft = np.fft.fft2(Y)
    dft_shift = np.fft.fftshift(dft)
    
    # 2. 在水印区域嵌入信息（高频区域）
    rows, cols = Y.shape
    watermark_mask = np.zeros_like(Y)
    
    # 在四个角落嵌入水印（高频区域）
    corner_size = 50
    corners = [
        (0, 0), (0, cols-corner_size), 
        (rows-corner_size, 0), (rows-corner_size, cols-corner_size)
    ]
    
    # 将文本转为二进制嵌入
    binary_text = ''.join(format(ord(c), '08b') for c in watermark_text)
    bit_idx = 0
    
    for corner in corners:
        y, x = corner
        for i in range(corner_size):
            for j in range(corner_size):
                if bit_idx < len(binary_text):
                    # 修改幅度谱的LSB
                    mag = np.abs(dft_shift[y+i, x+j])
                    modified_mag = mag - (mag % 2) + int(binary_text[bit_idx])
                    phase = np.angle(dft_shift[y+i, x+j])
                    dft_shift[y+i, x+j] = modified_mag * np.exp(1j*phase)
                    bit_idx += 1
    
    # 3. IDFT重建
    idft_shift = np.fft.ifftshift(dft_shift)
    Y_watermarked = np.fft.ifft2(idft_shift)
    yuv[:,:,0] = np.abs(Y_watermarked)
    watermarked_img = cv2.cvtColor(yuv, cv2.COLOR_YUV2BGR)
    
    return watermarked_img, binary_text

def detect_watermark(attacked_img, original_bits):
    # 提取处理后的图像Y通道
    yuv = cv2.cvtColor(attacked_img, cv2.COLOR_BGR2YUV)
    Y = yuv[:,:,0].astype(np.float32)
    
    # DFT变换
    dft = np.fft.fft2(Y)
    dft_shift = np.fft.fftshift(dft)
    
    # 检测水印
    detected_bits = ""
    rows, cols = Y.shape
    corner_size = 50
    corners = [
        (0, 0), (0, cols-corner_size), 
        (rows-corner_size, 0), (rows-corner_size, cols-corner_size)
    ]
    
    for corner in corners:
        y, x = corner
        for i in range(corner_size):
            for j in range(corner_size):
                if len(detected_bits) < len(original_bits):
                    mag = np.abs(dft_shift[y+i, x+j])
                    detected_bits += str(int(mag % 2))
    
    # 计算误码率
    error_rate = sum(1 for a,b in zip(original_bits, detected_bits) if a != b) / len(original_bits)
    return detected_bits, error_rate

# 模拟攻击函数
def simulate_attacks(img):
    # 1. 裁剪攻击（保留中心区域）
    h, w = img.shape[:2]
    cropped = img[h//4:3*h//4, w//4:3*w//4]
    
    # 2. 平移攻击（向右下平移）
    M = np.float32([[1,0,50],[0,1,50]])
    translated = cv2.warpAffine(img, M, (w,h))
    
    # 3. 缩放攻击（缩小到50%）
    scaled = cv2.resize(img, None, fx=0.5, fy=0.5)
    
    # 4. 涂抹攻击（高斯模糊）
    blurred = cv2.GaussianBlur(img, (15,15), 5)
    
    return cropped, translated, scaled, blurred

# 主流程
img_path = "Cameraman.bmp"
original_img = cv2.imread(img_path)
watermarked_img, original_bits = embed_watermark(img_path)

# 模拟攻击
cropped, translated, scaled, blurred = simulate_attacks(watermarked_img)

# 检测水印
attacks = {
    "Cropped": cropped,
    "Translated": translated, 
    "Scaled": scaled,
    "Blurred": blurred
}

# 创建更大的画布
plt.figure(figsize=(18, 12))

# 显示原图
plt.subplot(3, 3, 1)
plt.imshow(cv2.cvtColor(original_img, cv2.COLOR_BGR2RGB))
plt.title("Original Image")
plt.axis('off')

# 显示加水印后的图
plt.subplot(3, 3, 2)
plt.imshow(cv2.cvtColor(watermarked_img, cv2.COLOR_BGR2RGB))
plt.title("Watermarked Image")
plt.axis('off')

# 显示各种攻击后的结果
for i, (name, img) in enumerate(attacks.items()):
    detected_bits, error_rate = detect_watermark(img, original_bits)
    
    plt.subplot(3, 3, i+4)  # 调整位置到第三行
    plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    plt.title(f"{name} Attack\nError Rate: {error_rate:.2%}")
    plt.axis('off')

plt.tight_layout()
plt.show()